import pickle
import tarfile

import numpy as np

from implement.datasource.cifar10 import CIFAR10
from implement.transforms.astype import ToFloat
from implement.transforms.compose import Compose
from implement.transforms.normalize import Normalize
from utils.file_opera import get_file, load_cache_npz, save_cache_npz


class CIFAR100(CIFAR10):
    """
    CIFAR-100 数据集类，继承自 CIFAR-10。

    Args:
        train (bool): 是否为训练集。默认为 True。
        transform (Transform, optional): 数据转换操作。默认为 Compose([ToFloat(), Normalize(mean=0.5, std=0.5)])。
        target_transform (Transform, optional): 标签转换操作。默认为 None。
        label_type (str, optional): 标签类型，可选值为 'fine' 或 'coarse'。默认为 'fine'。
    """

    def __init__(self, train=True,
                 transform=Compose([ToFloat(), Normalize(mean=0.5, std=0.5)]),
                 target_transform=None,
                 label_type='fine'):
        """
        初始化 CIFAR-100 数据集。

        Args:
            train (bool): 是否为训练集。默认为 True。
            transform (Transform, optional): 数据转换操作。默认为 Compose([ToFloat(), Normalize(mean=0.5, std=0.5)])。
            target_transform (Transform, optional): 标签转换操作。默认为 None。
            label_type (str, optional): 标签类型，可选值为 'fine' 或 'coarse'。默认为 'fine'。
        """
        assert label_type in ['fine', 'coarse']
        self.label_type = label_type
        super().__init__(train, transform, target_transform)

    def prepare(self):
        """
        准备数据集。加载或下载数据。
        """
        url = 'https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz'
        self.data, self.label = load_cache_npz(url, self.train)
        if self.data is not None:
            return

        filepath = get_file(url)
        if self.train:
            self.data = self._load_data(filepath, 'train')
            self.label = self._load_label(filepath, 'train')
        else:
            self.data = self._load_data(filepath, 'test')
            self.label = self._load_label(filepath, 'test')
        self.data = self.data.reshape(-1, 3, 32, 32)
        save_cache_npz(self.data, self.label, url, self.train)

    def _load_data(self, filename, data_type='train', idx=None):
        """
        加载数据。

        Args:
            filename (str): 数据文件名。
            data_type (str): 数据类型，'train' 或 'test'。

        Returns:
            numpy.ndarray: 加载的数据。
        """
        with tarfile.open(filename, 'r:gz') as file:
            for item in file.getmembers():
                if data_type in item.name:
                    data_dict = pickle.load(file.extractfile(item), encoding='bytes')
                    data = data_dict[b'data']
                    return data

    def _load_label(self, filename, data_type='train', idx=None):
        """
        加载标签。

        Args:
            filename (str): 数据文件名。
            data_type (str): 数据类型，'train' 或 'test'。

        Returns:
            numpy.ndarray: 加载的标签。
        """
        assert data_type in ['train', 'test']
        with tarfile.open(filename, 'r:gz') as file:
            for item in file.getmembers():
                if data_type in item.name:
                    data_dict = pickle.load(file.extractfile(item), encoding='bytes')
                    if self.label_type == 'fine':
                        return np.array(data_dict[b'fine_labels'])
                    elif self.label_type == 'coarse':
                        return np.array(data_dict[b'coarse_labels'])

    @staticmethod
    def labels(label_type='fine'):
        """
        获取标签字典。

        Args:
            label_type (str, optional): 标签类型，'fine' 或 'coarse'。默认为 'fine'。

        Returns:
            dict: 标签字典。
        """
        coarse_labels = dict(enumerate(['aquatic mammals', 'fish', 'flowers', 'food containers', 'fruit and vegetables',
                                        'household electrical device', 'household furniture', 'insects',
                                        'large carnivores', 'large man-made outdoor things',
                                        'large natural outdoor scenes', 'large omnivores and herbivores',
                                        'medium-sized mammals', 'non-insect invertebrates', 'people', 'reptiles',
                                        'small mammals', 'trees', 'vehicles 1', 'vehicles 2']))
        fine_labels = []
        return fine_labels if label_type is 'fine' else coarse_labels
